// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Author: Markus Gutschke

#include <arpa/nameser.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <poll.h>
#include <resolv.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include "check.h"
#include "net.h"
#include "stringprintf.h"


static int openHostPort(const char *hostname, const char *service,
                        int port) {
  // Look up host and service
  char buf[6];
  if (port <= 0) {
    // Use build-in service definitions, where possible. This makes the
    // code work on embedded devices that might not have an /etc/services
    // database.
    check(service);
    if (!strcmp(service, "xmpp-client")) {
      service = "5222";
    } else if (!strcmp(service, "http")) {
      service = "80";
    } else if (!strcmp(service, "https")) {
      service = "443";
    } else if (!strcmp(service, "smtp")) {
      service = "25";
    } else if (!strcmp(service, "submission")) {
      service = "587";
    }
  } else {
    check(port > 0 && port < 65536);
    sprintf(buf, "%d", port);
    service = buf;
  }
  struct addrinfo hints = { .ai_family   = AF_UNSPEC,
                            .ai_socktype = SOCK_STREAM,
                            .ai_protocol = IPPROTO_TCP,
                            .ai_flags    = AI_ADDRCONFIG | AI_V4MAPPED
                          };
 family:;
  struct addrinfo *server_info = NULL;
  if (getaddrinfo(hostname, service, &hints, &server_info) < 0) {
    // uCLibc can't handle ai_family == AF_UNSPEC. So, we instead explicitly
    // try both AF_INET and AF_INET6.
  error:
    switch (hints.ai_family) {
    case AF_UNSPEC:
      hints.ai_family = AF_INET6;
      goto family;
    case AF_INET6:
      hints.ai_family = AF_INET;
      goto family;
    default:
      break;
    }
    return -1;
  }

  // Find first address info that actually can be connected to.
  int fd = -1;
  for (struct addrinfo *info = server_info; info; info = info->ai_next) {
    fd = socket(info->ai_family, info->ai_socktype, info->ai_protocol);
    if (fd < 0) {
      continue;
    }
    // Make our socket non-blocking, so that we can read data even if we
    // don't know how many bytes are available for reading.
    check(!fcntl(fd, F_SETFL, (long)O_NONBLOCK));

    // Set send and receive time outs. If we would otherwise block, we'd
    // rather be waiting in poll()
    struct timeval tmo = { .tv_sec = 5 };
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tmo, sizeof(tmo));
    setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tmo, sizeof(tmo));

    if (connect(fd, info->ai_addr, info->ai_addrlen) < 0 &&
        errno != EINPROGRESS) {
    try_next_address:
      close(fd);
      fd = -1;
      continue;
    }
    
    struct pollfd pfd = {
      .fd     = fd,
      .events = POLLOUT
    };
    if (TEMP_FAILURE_RETRY(poll(&pfd, 1, 5000)) != 1 ||
        pfd.revents != POLLOUT) {
      goto try_next_address;
    }

    break;
  }
  freeaddrinfo(server_info);

  if (fd == -1) {
    // Maybe, there are other address families that we can still try.
    goto error;
  }

  return fd;
}

struct SRV {
  const char *hostname;
  int        port;
  int        pref;
  int        weight;
};

static int cmpSRV(const void *a, const void *b) {
  const struct SRV *as = (const struct SRV *)a;
  const struct SRV *bs = (const struct SRV *)b;
  return as->pref == bs->pref ? bs->weight - as->weight
                              : as->pref   - bs->pref;
}

static int openService(const char *hostname, const char *service) {
  // Assemble SRV style hostname, or use plain hostname for MX lookups
  int ns_type = strcmp(service, "smtp") ? ns_t_srv : ns_t_mx;
  char *name;
  if (ns_type == ns_t_srv) {
    name = stringPrintf(NULL, "_%s._tcp.%s", service, hostname);
  } else {
    name = strdup(hostname);
  }

  // Query DNS server
  union {
    HEADER hdr;
    u_char buf[1500];
  } resp;
  ssize_t sz;
  struct SRV *srv = NULL;
  ssize_t nsrv = 0;
  int fd = -1;
  if (res_init() < 0 ||
      (sz = res_query(name, ns_c_in, ns_type, resp.buf,
                      sizeof(resp))) < sizeof(resp.hdr)) {
    free(name);
    goto done;
  }
  free(name);

  // Skip query record
  union {
    u_char   *c;
    uint16_t *u16;
    uint32_t *u32;
  } ptr = { .c = &resp.buf[sizeof(HEADER)] };
  char buf[256];
  for (int queries = ntohs(resp.hdr.qdcount); queries--; ) {
    ssize_t len = dn_expand(resp.buf, &resp.buf[sz], ptr.c, buf, sizeof(buf));
    if (len < 0 || ptr.c + len + QFIXEDSZ >= &resp.buf[sz]) {
      goto done;
    }
    ptr.c += len + QFIXEDSZ;
  }
  
  // Parse answers
  check(srv = malloc(sizeof(struct SRV) * ntohs(resp.hdr.ancount)));
  memset(srv, 0, sizeof(struct SRV) * ntohs(resp.hdr.ancount));
  for (int answers = ntohs(resp.hdr.ancount); answers--; ) {
    ssize_t len = dn_expand(resp.buf, &resp.buf[sz], ptr.c, buf, sizeof(buf));
    if (len < 0 || ((ptr.c += len) + 10) > &resp.buf[sz]) {
      goto done;
    }
    int type = ntohs(*ptr.u16++);
    /* class */       ptr.u16++;
    /* ttl   */       ptr.u32++;
    int dlen = ntohs(*ptr.u16++);
    if (type != (ns_type == ns_t_srv ? T_SRV : T_MX)) {
      if ((ptr.c += dlen) > &resp.buf[sz]) {
        goto done;
      }
      continue;
    } else if (ptr.c + (ns_type == ns_t_srv ? 6 : 2) > &resp.buf[sz]) {
      goto done;
    }
    // We treat MX records as if they were a subclass of SRV records
    srv[nsrv].pref   = ntohs(*ptr.u16++);
    srv[nsrv].weight = (ns_type == ns_t_srv ? ntohs(*ptr.u16++) : 1) *
                       (rand() & 0x7FFF);
    srv[nsrv].port   = (ns_type == ns_t_srv ? ntohs(*ptr.u16++) : 25);
    if ((len = dn_expand(resp.buf, &resp.buf[sz], ptr.c, buf, sizeof(buf)))<0||
        (ptr.c += len) > &resp.buf[sz]) {
      goto done;
    }
    check(srv[nsrv++].hostname = strdup(buf));
  }

  // Sort by preference and weight
  qsort(srv, nsrv, sizeof(struct SRV), cmpSRV);

  for (int i = 0; i < nsrv; ++i) {
    if ((fd = openHostPort(srv[i].hostname, service, srv[i].port)) >= 0) {
      break;
    }
  }

  if (fd < 0 && ns_type == ns_t_mx) {
    // Try the submission port in addition to the smtp port. Port 25 might be
    // blocked by the firewall.
    for (int i = 0; i < nsrv; ++i) {
      if ((fd = openHostPort(srv[i].hostname, service, 587)) >= 0) {
        break;
      }
    }
  }

  // Clean up temporary data structures
 done:
  for (int i = 0; i < nsrv; ++i) {
    free((void *)srv[i].hostname);
  }
  free(srv);
  return fd;
}

int openSocket(const char *hostname, const char *service) {
  // Look up SRV or MX record first
  int fd = openService(hostname, service);
  if (fd < 0) {
    // If no SRV or MX record could be found, try connecting to the named
    // server on the well-known service port.
    fd = openHostPort(hostname, service, -1);
    if (fd < 0 && !strcmp(service, "smtp")) {
      fd = openHostPort(hostname, "submission", -1);
    }
  }
  if (fd < 0) {
    return -1;
  }
  return fd;
}
